package com.mm.ai.service.impl;

import com.fasterxml.jackson.databind.ObjectMapper;
import com.google.common.collect.Lists;
import com.mm.ai.common.CommonException;
import com.mm.ai.common.constant.CozeApiConstant;
import com.mm.ai.dto.coze.CozeBaseRspDTO;
import com.mm.ai.dto.coze.CreateknowledgeReqDTO;
import com.mm.ai.dto.coze.CreateknowledgeRspDTO;
import com.mm.ai.service.CozeKnowledgeService;
import com.mm.ai.utils.RestHttpUtil;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang.StringUtils;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.http.HttpEntity;
import org.springframework.http.HttpHeaders;
import org.springframework.http.MediaType;
import org.springframework.stereotype.Service;
import org.springframework.web.client.RestTemplate;
import org.springframework.web.multipart.MultipartFile;

import java.io.IOException;
import java.util.Base64;
import java.util.HashMap;
import java.util.Set;

/**
 * @Author：tangnaichao
 * @Date：2024/10/18 19:51
 */
@Service
@Slf4j
public class CozeKnowledgeServiceImpl implements CozeKnowledgeService {
    @Autowired
    private RestHttpUtil restHttpUtil;
    @Autowired
    private RestTemplate restTemplate;

    @Override
    public CreateknowledgeRspDTO createKnowledge(MultipartFile file) {
        byte[] fileContent = new byte[0];
        try {
            fileContent = file.getBytes();
        } catch (IOException e) {
            e.printStackTrace();
        }
        String base64Encoded = Base64.getEncoder().encodeToString(fileContent);

        // 构造请求体
        CreateknowledgeReqDTO reqDTO = new CreateknowledgeReqDTO();
        reqDTO.setDataset_id(7427316483476471858L); // 替换为您的数据集 ID
        CreateknowledgeReqDTO.DocumentBase db = new CreateknowledgeReqDTO.DocumentBase();
        db.setName(file.getOriginalFilename());

        CreateknowledgeReqDTO.SourceInfo si = new CreateknowledgeReqDTO.SourceInfo();
        si.setFile_base64(base64Encoded);
        si.setFile_type("txt"); // 或者指定一个默认的文件类型，如 "txt"
        db.setSource_info(si);
        reqDTO.setDocument_bases(Lists.newArrayList(db));

        String knowledgeList= selectListByKnowledgeBaseIdPage(7427316483476471858L,1,10);
        if(StringUtils.isEmpty(knowledgeList)){
            CreateknowledgeReqDTO.ChunkStrategy chunkStrategy = new CreateknowledgeReqDTO.ChunkStrategy();
            chunkStrategy.setSeparator("\n\n");
            chunkStrategy.setMax_tokens(800);
            chunkStrategy.setRemove_extra_spaces(false);
            chunkStrategy.setRemove_urls_emails(false);
            chunkStrategy.setChunk_type(1);
            reqDTO.setChunk_strategy(chunkStrategy);
        }

        // 序列化请求对象
        String requestJson = null;
        try {
            requestJson = new ObjectMapper().writeValueAsString(reqDTO);
        } catch (Exception e) {
            throw new CommonException("Error sending request: " + e.getMessage());
        }

        // 设置请求头
        HttpHeaders headers = new HttpHeaders();
        headers.setContentType(MediaType.APPLICATION_JSON);
        headers.set("Authorization", CozeApiConstant.TONGXINHUI_APITOKEN);
        headers.set("Agw-Js-Conv", "str");

        HttpEntity<String> entity = new HttpEntity<>(requestJson, headers);

        // 发送请求
        try {
            CreateknowledgeRspDTO response = restTemplate.postForObject(CozeApiConstant.CREATE_KNOWLEDGE_COZE_API_URL_PRO, entity, CreateknowledgeRspDTO.class);
            return response;
        } catch (Exception e) {
            throw new CommonException("Error sending request: " + e.getMessage());
        }
    }

    @Override
    public void updateknowledgeById(Long documentId, String documentName) {
        String requestJson = null;
        try {
            requestJson = new ObjectMapper().writeValueAsString(new HashMap() {{
                put("document_id", documentId);
                put("document_name", documentName);
            }});
        } catch (Exception e) {
            throw new CommonException("Error sending request: " + e.getMessage());
        }

        // 设置请求头
        HttpHeaders headers = new HttpHeaders();
        headers.setContentType(MediaType.APPLICATION_JSON);
        headers.set("Authorization", CozeApiConstant.TONGXINHUI_APITOKEN);
        headers.set("Agw-Js-Conv", "str");

        HttpEntity<String> entity = new HttpEntity<>(requestJson, headers);

        // 发送请求
        try {
            CozeBaseRspDTO response = restTemplate.postForObject(CozeApiConstant.UPDATE_KNOWLEDGE_COZE_API_URL_PRO, entity, CozeBaseRspDTO.class);
        } catch (Exception e) {
            throw new CommonException("Error sending request: " + e.getMessage());
        }
    }

    @Override
    public void deleteByDocumentIdBatch(Set<String> documentIds) {
        String requestJson = null;
        try {
            requestJson = new ObjectMapper().writeValueAsString(new HashMap() {{
                put("document_ids", documentIds);
            }});
        } catch (Exception e) {
            throw new RuntimeException("Error serializing request: " + e.getMessage(), e);
        }

        // 设置请求头
        HttpHeaders headers = new HttpHeaders();
        headers.setContentType(MediaType.APPLICATION_JSON);
        headers.set("Authorization", CozeApiConstant.TONGXINHUI_APITOKEN);
        headers.set("Agw-Js-Conv", "str");

        HttpEntity<String> entity = new HttpEntity<>(requestJson, headers);

        // 发送请求
        try {
            CozeBaseRspDTO response = restTemplate.postForObject(CozeApiConstant.DELETE_KNOWLEDGE_COZE_API_URL_PRO, entity, CozeBaseRspDTO.class);
        } catch (Exception e) {
            throw new CommonException("Error sending request: " + e.getMessage());
        }
    }

    @Override
    public String selectListByKnowledgeBaseIdPage(Long datasetId, Integer currentPage, Integer pageSize) {
        String requestJson = null;
        try {
            requestJson = new ObjectMapper().writeValueAsString(new HashMap() {{
                put("dataset_id", datasetId);
                put("page", currentPage);
                put("size", currentPage);
            }});
        } catch (Exception e) {
            throw new CommonException("Error sending request: " + e.getMessage());
        }

        // 设置请求头
        HttpHeaders headers = new HttpHeaders();
        headers.setContentType(MediaType.APPLICATION_JSON);
        headers.set("Authorization", CozeApiConstant.TONGXINHUI_APITOKEN);
        headers.set("Agw-Js-Conv", "str");

        HttpEntity<String> entity = new HttpEntity<>(requestJson, headers);

        // 发送请求
        try {
            return restTemplate.postForObject(CozeApiConstant.LIST_KNOWLEDGE_COZE_API_URL_PRO, entity, String.class);
        } catch (Exception e) {
            throw new CommonException("Error sending request: " + e.getMessage());
        }
    }
}
